import base64
import os.path
import time

from flask import Flask
from flask import request
from utils import gen_uuid
from utils.logtools import logger
from utils.audio.tools import bytes2tensor
from lid.lid_proccesser import WenetEncoder
from lid.lid_proccesser import MyClassifier

app = Flask(__name__)


def init_classifier():
    model_name = "langid-ecapa-encoder-512-epoch-5"
    basedir = os.path.dirname(__file__)
    lid_model_dir = os.path.join(basedir, f"lid/models/{model_name}")
    encoder_path = os.path.join(basedir, f"lid/models/encoder.onnx")

    encoder = WenetEncoder(encoder_path=encoder_path)
    classifier = MyClassifier.from_hparams(source=lid_model_dir, savedir=lid_model_dir)
    classifier.set_encoder(encoder)
    logger.info("lid_module innit success")
    return classifier


classifier = init_classifier()

@app.route('/')
def hello_world():
    return 'Hello World!'


@app.route('/_health')
def health():
    return {"status_code": 200}, 200


@app.route('/speech-language-detection', methods=['POST'])
def speech_language_detection():
    params = request.get_json()
    session_id = params.get("session_id", gen_uuid())
    pcm_base64 = params.get('data', None)
    if pcm_base64 is None:
        return {
            "status_code": 400,
            "task_id": session_id,
            "err_msg": "audio data can't be empty"
        }, 400

    lid = 1

    try:
        bytes_to_decode = pcm_base64.encode("utf-8")
        pcm = base64.b64decode(bytes_to_decode)
        logger.info({
            "sesssion_id": session_id,
            "status": "revived stream",
            "stream length": len(pcm)
        })
        waveform = bytes2tensor(pcm)
        t0 = time.time()
        lid = classifier.detect_sentenc_lid(waveform)
        t1 = time.time()
        logger.info({
            "session_id": session_id,
            "status": "finish_detect",
            "time_cost": "{:.2f}s".format(t1 - t0)
        })
        return {"status_code": 200, "task_id": session_id, "language_id": lid}, 200

    except Exception as ex:
        logger.error({"task_id": session_id, "status": "task failed"})
        logger.error(ex)
        return {
            "status_code": 500,
            "task_id": session_id,
            "status": "internal error"
        }, 500
    return {"status_code": 200, "task_id": session_id, "language_id": lid}, 200


if __name__ == '__main__':
    app.run(host="0.0.0.0", port=8000)
